import onnx
from onnx.shape_inference import infer_shapes
from onnx import load_model, save_model
import argparse

parser = argparse.ArgumentParser(description="onnx infer shape tool.")
parser.add_argument('input_model_file', type=str, help='input model file')
parser.add_argument('out_model_file', type=str , help='out model file.')

args = parser.parse_args()

onnx_model = onnx.load(args.input_model_file) 
onnx_model = infer_shapes(onnx_model) 
save_model(onnx_model, args.out_model_file)

